import zstandard as zstd  # pip install zstandard

# 转为二进制字符串：将字节流还原为 '0' 和 '1' 组成的字符串
def bytes_to_bin_str(data: bytes) -> str:
    return ''.join(f'{byte:08b}' for byte in data)

def compressMethod(bits: list[str], level=6)->bytes:
    """
    使用 Zstandard (zstd) 压缩二进制序列
    :param lsbs: 字符 '0' 和 '1' 组成的列表
    :param level: 压缩级别 (1-22)，默认 6（推荐）
    :return: 压缩后的字节数据
    """
    lsbsStr = ''.join(bits)
    
    # 转为字节：每8位一组
    def bin_str_to_bytes(s):
        padded = s.ljust((len(s) + 7) // 8 * 8, '0')  # 补齐到8的倍数
        return bytes([int(padded[i:i+8], 2) for i in range(0, len(padded), 8)])
    
    data = bin_str_to_bytes(lsbsStr)
    original_size = len(data)

    # 创建压缩器
    cctx = zstd.ZstdCompressor(level=level)
    compressed = cctx.compress(data)

    # print(f"Zstd 压缩前: {original_size} 字节")
    # print(f"Zstd 压缩后: {len(compressed)} 字节")
    # ratio = len(compressed) / original_size if original_size > 0 else 0
    # print(f"压缩比: {ratio:.2f}")
    # print(f"压缩级别: {level}")

    return compressed

# 解压：将压缩数据byte解压还原为bit列表（list[str]）
def decompressMethod(compressed_data: bytes)-> list[str]:
    """
    解压 zstd 压缩的数据
    :param compressed_data: 压缩后的字节数据
    :return: bytes
    """
    # 解压
    dctx = zstd.ZstdDecompressor()
    decompressed_bytes = dctx.decompress(compressed_data)
    bin_str = bytes_to_bin_str(decompressed_bytes)
    return list(bin_str)

# === 示例使用 ===
if __name__ == "__main__":
    # 示例数据：模拟有结构的二进制序列
    bits = '10' * 500 
    bits = list(bits)
    print("=== Zstd 压缩测试 ===")
    compressed = compressMethod(bits, level=6)

    print("\n=== Zstd 解压测试 ===")
    decompressed = decompressMethod(compressed)

    # 验证
    if decompressed == bits:
        print("解压成功，数据一致！")
    else:
        print("解压失败！")